#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""展示如何让backtrader的绘图具有动态缩放功能"""

import backtrader as bt
import matplotlib.pyplot as plt
import datetime
import pandas as pd
from backtrader import Plot_OldSync, Plot

class SimpleMovingAverageStrategy(bt.Strategy):
    """简单移动平均线策略用于演示"""
    params = (
        ('maperiod', 15),
    )

    def __init__(self):
        self.dataclose = self.datas[0].close
        # 添加移动平均线指标
        self.sma = bt.indicators.SimpleMovingAverage(
            self.datas[0], period=self.params.maperiod)

    def next(self):
        # 简单的策略逻辑：价格上穿均线买入，下穿均线卖出
        if not self.position:
            if self.dataclose[0] > self.sma[0]:
                self.buy()
        else:
            if self.dataclose[0] < self.sma[0]:
                self.sell()

# 方法1: 使用matplotlib的默认交互模式
def plot_with_matplotlib_interactive():
    print("方法1: 使用matplotlib的默认交互模式")
    cerebro = bt.Cerebro()
    
    # 添加策略
    cerebro.addstrategy(SimpleMovingAverageStrategy)
    
    # 创建数据
    data = bt.feeds.YahooFinanceData(
        dataname='AAPL',
        fromdate=datetime.datetime(2020, 1, 1),
        todate=datetime.datetime(2021, 12, 31)
    )
    
    cerebro.adddata(data)
    cerebro.broker.setcash(100000.0)
    cerebro.addsizer(bt.sizers.FixedSize, stake=10)
    cerebro.broker.setcommission(commission=0.001)
    
    # 运行策略
    print('初始资金: %.2f' % cerebro.broker.getvalue())
    cerebro.run()
    print('最终资金: %.2f' % cerebro.broker.getvalue())
    
    # 使用matplotlib的交互模式绘图
    # iplot=False表示不在notebook中内联显示，而是弹出独立窗口
    # matplotlib默认窗口支持缩放、平移等交互功能
    plt.ion()  # 启用交互模式
    cerebro.plot(iplot=False, style='candle')
    plt.ioff()  # 关闭交互模式
    plt.show()  # 显示图表

# 方法2: 使用Plot_OldSync绘图器
def plot_with_old_sync():
    print("方法2: 使用Plot_OldSync绘图器")
    cerebro = bt.Cerebro()
    
    # 添加策略
    cerebro.addstrategy(SimpleMovingAverageStrategy)
    
    # 创建数据
    data = bt.feeds.YahooFinanceData(
        dataname='MSFT',
        fromdate=datetime.datetime(2020, 1, 1),
        todate=datetime.datetime(2021, 12, 31)
    )
    
    cerebro.adddata(data)
    cerebro.broker.setcash(100000.0)
    cerebro.addsizer(bt.sizers.FixedSize, stake=10)
    cerebro.broker.setcommission(commission=0.001)
    
    # 运行策略
    cerebro.run()
    
    # 使用Plot_OldSync绘图器，它提供了更灵活的图表控制
    plotter = Plot_OldSync(
        style='candle',
        barup='red',
        bardown='green',
        volup='#ffcccc',
        voldown='#ccffcc',
    )
    
    cerebro.plot(plotter=plotter, iplot=False)

# 方法3: 创建自定义plotter类，扩展交互功能
class InteractivePlotter(Plot):
    """自定义绘图器，增强交互功能"""
    def __init__(self, **kwargs):
        super(InteractivePlotter, self).__init__(**kwargs)
        # 可以在这里设置额外的交互选项
    
    def show(self):
        """重写show方法，增强交互体验"""
        for f in self.figs:
            # 设置窗口标题
            f.canvas.manager.set_window_title('Backtrader Interactive Chart')
            
            # 设置鼠标中键拖动平移
            f.canvas.toolbar.pan()
            
            # 可以添加额外的交互设置
            # 例如：添加十字光标
            # 注意：backtrader内部已经实现了MultiCursor
        
        # 显示图表
        plt.show()

def plot_with_custom_plotter():
    print("方法3: 使用自定义绘图器增强交互功能")
    cerebro = bt.Cerebro()
    
    # 添加策略
    cerebro.addstrategy(SimpleMovingAverageStrategy)
    
    # 创建数据
    data = bt.feeds.YahooFinanceData(
        dataname='GOOG',
        fromdate=datetime.datetime(2020, 1, 1),
        todate=datetime.datetime(2021, 12, 31)
    )
    
    cerebro.adddata(data)
    cerebro.broker.setcash(100000.0)
    cerebro.addsizer(bt.sizers.FixedSize, stake=10)
    cerebro.broker.setcommission(commission=0.001)
    
    # 运行策略
    cerebro.run()
    
    # 使用自定义绘图器
    custom_plotter = InteractivePlotter(
        style='candle',
        barup='red',
        bardown='green',
    )
    
    cerebro.plot(plotter=custom_plotter, iplot=False)

# 方法4: 使用plotly创建交互式网页图表
def plot_with_plotly():
    """使用plotly创建高度交互式的图表"""
    print("方法4: 使用plotly创建交互式网页图表")
    
    try:
        import plotly.graph_objects as go
        import plotly.offline as pyo
        
        # 创建数据
        data = bt.feeds.YahooFinanceData(
            dataname='TSLA',
            fromdate=datetime.datetime(2020, 1, 1),
            todate=datetime.datetime(2021, 12, 31),
            reverse=False
        )
        
        # 提取数据到pandas DataFrame
        # 注意：这只是一个示例，实际使用时需要运行策略并收集数据
        df = pd.DataFrame({
            'Date': pd.date_range('2020-01-01', periods=500),
            'Open': range(500),
            'High': [x + 10 for x in range(500)],
            'Low': [x - 10 for x in range(500)],
            'Close': range(500),
            'Volume': [x * 100 for x in range(500)]
        })
        
        # 创建蜡烛图
        fig = go.Figure(data=[
            go.Candlestick(
                x=df['Date'],
                open=df['Open'],
                high=df['High'],
                low=df['Low'],
                close=df['Close'],
                name='价格'
            ),
            go.Scatter(
                x=df['Date'],
                y=df['Close'].rolling(window=15).mean(),
                mode='lines',
                name='15日均线',
                line=dict(color='blue')
            )
        ])
        
        # 添加成交量子图
        fig.add_trace(
            go.Bar(
                x=df['Date'],
                y=df['Volume'],
                name='成交量',
                yaxis='y2'
            )
        )
        
        # 更新布局以支持缩放
        fig.update_layout(
            title='交互式交易图表 (支持动态缩放)',
            yaxis_title='价格',
            yaxis2=dict(
                title='成交量',
                overlaying='y',
                side='right'
            ),
            xaxis=dict(
                rangeslider=dict(visible=True),  # 底部时间范围滑块
                type='date'
            ),
            # 添加交互功能配置
            hovermode='x unified',
            dragmode='zoom'
        )
        
        # 保存为HTML文件或在浏览器中显示
        pyo.plot(fig, filename='interactive_chart.html', auto_open=True)
        print("交互式图表已在浏览器中打开")
    except ImportError:
        print("请安装plotly库: pip install plotly")
    except Exception as e:
        print(f"创建plotly图表时出错: {e}")

if __name__ == '__main__':
    print("backtrader交互式绘图示例")
    print("="*50)
    
    # 选择要使用的方法
    # 方法1: matplotlib默认交互模式
    # plot_with_matplotlib_interactive()
    
    # 方法2: 使用Plot_OldSync
    # plot_with_old_sync()
    
    # 方法3: 自定义绘图器
    # plot_with_custom_plotter()
    
    # 方法4: 使用plotly创建交互式网页图表
    plot_with_plotly()